"""
Database Manager
Centralized database connection and transaction management
"""

import logging
import asyncio
from contextlib import asynccontextmanager
from typing import Dict, List, Optional, Any, AsyncGenerator
import asyncpg
import psycopg2
from psycopg2.extras import RealDictCursor
from psycopg2.pool import ThreadedConnectionPool

from config import get_config

logger = logging.getLogger(__name__)

class DatabaseManager:
    """
    Centralized database connection management with connection pooling
    """
    
    def __init__(self):
        self.config = get_config()
        self.pool = None
        self.async_pool = None
        self._initialize_connection_pool()
    
    def _initialize_connection_pool(self):
        """Initialize database connection pool"""
        try:
            # Parse database URL
            db_url = self.config.database.url
            
            # Create synchronous connection pool for legacy compatibility
            self.pool = ThreadedConnectionPool(
                minconn=1,
                maxconn=self.config.database.pool_size,
                dsn=db_url,
                cursor_factory=RealDictCursor
            )
            
            logger.info(f"Database connection pool initialized with {self.config.database.pool_size} connections")
            
        except Exception as e:
            logger.error(f"Failed to initialize database connection pool: {e}")
            self.pool = None
    
    async def _initialize_async_pool(self):
        """Initialize async database connection pool"""
        if self.async_pool is None:
            try:
                db_url = self.config.database.url
                self.async_pool = await asyncpg.create_pool(
                    db_url,
                    min_size=1,
                    max_size=self.config.database.pool_size,
                    command_timeout=30
                )
                logger.info("Async database connection pool initialized")
            except Exception as e:
                logger.error(f"Failed to initialize async database pool: {e}")
                raise
    
    @asynccontextmanager
    async def get_connection(self) -> AsyncGenerator[asyncpg.Connection, None]:
        """Get async database connection"""
        if self.async_pool is None:
            await self._initialize_async_pool()
        
        async with self.async_pool.acquire() as connection:
            yield connection
    
    def get_sync_connection(self):
        """Get synchronous database connection"""
        if self.pool is None:
            self._initialize_connection_pool()
        
        return self.pool.getconn()
    
    def return_sync_connection(self, conn):
        """Return synchronous database connection to pool"""
        if self.pool:
            self.pool.putconn(conn)
    
    async def execute_query(
        self, 
        query: str, 
        params: tuple = None,
        fetch: str = 'none'  # 'none', 'one', 'all'
    ) -> Any:
        """Execute a database query"""
        try:
            async with self.get_connection() as conn:
                if fetch == 'one':
                    return await conn.fetchrow(query, *(params or ()))
                elif fetch == 'all':
                    return await conn.fetch(query, *(params or ()))
                else:
                    return await conn.execute(query, *(params or ()))
        except Exception as e:
            logger.error(f"Database query failed: {e}")
            raise
    
    async def execute_transaction(self, queries: List[Dict[str, Any]]) -> bool:
        """Execute multiple queries in a transaction"""
        try:
            async with self.get_connection() as conn:
                async with conn.transaction():
                    for query_info in queries:
                        query = query_info['query']
                        params = query_info.get('params', ())
                        await conn.execute(query, *params)
            return True
        except Exception as e:
            logger.error(f"Database transaction failed: {e}")
            return False
    
    async def check_connection(self) -> bool:
        """Check if database connection is working"""
        try:
            result = await self.execute_query("SELECT 1", fetch='one')
            return result is not None
        except Exception:
            return False
    
    async def get_table_info(self, table_name: str) -> List[Dict[str, Any]]:
        """Get information about table columns"""
        try:
            query = """
                SELECT column_name, data_type, is_nullable, column_default
                FROM information_schema.columns
                WHERE table_name = $1
                ORDER BY ordinal_position
            """
            return await self.execute_query(query, (table_name,), fetch='all')
        except Exception as e:
            logger.error(f"Failed to get table info for {table_name}: {e}")
            return []
    
    async def create_table_if_not_exists(self, table_name: str, schema: str):
        """Create table if it doesn't exist"""
        try:
            await self.execute_query(f"CREATE TABLE IF NOT EXISTS {table_name} {schema}")
            logger.info(f"Table {table_name} created or verified")
        except Exception as e:
            logger.error(f"Failed to create table {table_name}: {e}")
            raise
    
    async def insert_or_update(
        self, 
        table: str, 
        data: Dict[str, Any], 
        conflict_columns: List[str],
        update_columns: Optional[List[str]] = None
    ) -> bool:
        """Insert or update record"""
        try:
            columns = list(data.keys())
            values = list(data.values())
            placeholders = [f"${i+1}" for i in range(len(values))]
            
            # Build conflict resolution
            conflict_str = ", ".join(conflict_columns)
            
            if update_columns:
                update_str = ", ".join([f"{col} = EXCLUDED.{col}" for col in update_columns])
                on_conflict = f"ON CONFLICT ({conflict_str}) DO UPDATE SET {update_str}"
            else:
                on_conflict = f"ON CONFLICT ({conflict_str}) DO NOTHING"
            
            query = f"""
                INSERT INTO {table} ({", ".join(columns)})
                VALUES ({", ".join(placeholders)})
                {on_conflict}
            """
            
            await self.execute_query(query, tuple(values))
            return True
            
        except Exception as e:
            logger.error(f"Insert/update failed for table {table}: {e}")
            return False
    
    async def get_by_id(self, table: str, id_value: Any, id_column: str = 'id') -> Optional[Dict[str, Any]]:
        """Get record by ID"""
        try:
            query = f"SELECT * FROM {table} WHERE {id_column} = $1"
            result = await self.execute_query(query, (id_value,), fetch='one')
            return dict(result) if result else None
        except Exception as e:
            logger.error(f"Failed to get record from {table} with {id_column}={id_value}: {e}")
            return None
    
    async def get_many(
        self, 
        table: str, 
        where_conditions: Dict[str, Any] = None,
        order_by: str = None,
        limit: int = None,
        offset: int = None
    ) -> List[Dict[str, Any]]:
        """Get multiple records with filtering"""
        try:
            query = f"SELECT * FROM {table}"
            params = []
            
            if where_conditions:
                where_clauses = []
                param_count = 1
                for column, value in where_conditions.items():
                    where_clauses.append(f"{column} = ${param_count}")
                    params.append(value)
                    param_count += 1
                query += " WHERE " + " AND ".join(where_clauses)
            
            if order_by:
                query += f" ORDER BY {order_by}"
            
            if limit:
                query += f" LIMIT ${len(params) + 1}"
                params.append(limit)
            
            if offset:
                query += f" OFFSET ${len(params) + 1}"
                params.append(offset)
            
            results = await self.execute_query(query, tuple(params), fetch='all')
            return [dict(row) for row in results]
            
        except Exception as e:
            logger.error(f"Failed to get records from {table}: {e}")
            return []
    
    async def count_records(
        self, 
        table: str, 
        where_conditions: Dict[str, Any] = None
    ) -> int:
        """Count records in table"""
        try:
            query = f"SELECT COUNT(*) FROM {table}"
            params = []
            
            if where_conditions:
                where_clauses = []
                param_count = 1
                for column, value in where_conditions.items():
                    where_clauses.append(f"{column} = ${param_count}")
                    params.append(value)
                    param_count += 1
                query += " WHERE " + " AND ".join(where_clauses)
            
            result = await self.execute_query(query, tuple(params), fetch='one')
            return result['count'] if result else 0
            
        except Exception as e:
            logger.error(f"Failed to count records in {table}: {e}")
            return 0
    
    async def delete_records(
        self, 
        table: str, 
        where_conditions: Dict[str, Any]
    ) -> int:
        """Delete records and return count of deleted rows"""
        try:
            where_clauses = []
            params = []
            param_count = 1
            
            for column, value in where_conditions.items():
                where_clauses.append(f"{column} = ${param_count}")
                params.append(value)
                param_count += 1
            
            query = f"DELETE FROM {table} WHERE " + " AND ".join(where_clauses)
            result = await self.execute_query(query, tuple(params))
            
            # Extract row count from result
            if hasattr(result, 'split'):
                # Result is like "DELETE 5"
                return int(result.split()[-1]) if result.split()[-1].isdigit() else 0
            return 0
            
        except Exception as e:
            logger.error(f"Failed to delete records from {table}: {e}")
            return 0
    
    async def update_records(
        self, 
        table: str, 
        updates: Dict[str, Any],
        where_conditions: Dict[str, Any]
    ) -> int:
        """Update records and return count of updated rows"""
        try:
            set_clauses = []
            where_clauses = []
            params = []
            param_count = 1
            
            # Build SET clause
            for column, value in updates.items():
                set_clauses.append(f"{column} = ${param_count}")
                params.append(value)
                param_count += 1
            
            # Build WHERE clause
            for column, value in where_conditions.items():
                where_clauses.append(f"{column} = ${param_count}")
                params.append(value)
                param_count += 1
            
            query = f"""
                UPDATE {table} 
                SET {', '.join(set_clauses)}
                WHERE {' AND '.join(where_clauses)}
            """
            
            result = await self.execute_query(query, tuple(params))
            
            # Extract row count from result
            if hasattr(result, 'split'):
                return int(result.split()[-1]) if result.split()[-1].isdigit() else 0
            return 0
            
        except Exception as e:
            logger.error(f"Failed to update records in {table}: {e}")
            return 0
    
    async def close_connections(self):
        """Close all database connections"""
        try:
            if self.async_pool:
                await self.async_pool.close()
                self.async_pool = None
            
            if self.pool:
                self.pool.closeall()
                self.pool = None
            
            logger.info("Database connections closed")
            
        except Exception as e:
            logger.error(f"Error closing database connections: {e}")
    
    async def get_database_stats(self) -> Dict[str, Any]:
        """Get database statistics"""
        try:
            stats = {}
            
            # Get table sizes
            table_stats_query = """
                SELECT 
                    schemaname,
                    tablename,
                    attname,
                    n_distinct,
                    avg_width
                FROM pg_stats 
                WHERE schemaname = 'public'
                ORDER BY tablename, attname
            """
            
            table_stats = await self.execute_query(table_stats_query, fetch='all')
            stats['table_stats'] = [dict(row) for row in table_stats]
            
            # Get connection info
            if self.async_pool:
                stats['connection_pool'] = {
                    'size': self.async_pool.get_size(),
                    'min_size': self.async_pool.get_min_size(),
                    'max_size': self.async_pool.get_max_size()
                }
            
            return stats
            
        except Exception as e:
            logger.error(f"Failed to get database stats: {e}")
            return {}

# Global database manager instance
db_manager = DatabaseManager()

# Convenience functions
def get_db_manager() -> DatabaseManager:
    """Get the global database manager instance"""
    return db_manager